#
# models.py
#

import numpy as np

np_fromiter = np.fromiter
np_zeros = np.zeros
np_empty = np.empty

class Model:
    #
    def evaluate(self, Xk):
        pass
    #
    def evaluate_all(self, X):
        evaluate = self.evaluate
        iter_eval = (evaluate(Xk) for Xk in X)
        return np.fromiter(iter_eval, 'd', len(X))        
    #
    def gradient(self, Xk):
        pass
    #
    def gradient_x(self, Xk):
        pass
    #
    def gradient_all(self, X):
        pass

class LinearModel(Model):
    #
    def __init__(self, n, param=None):
        self.n_input = n
        self.n_param = n+1
        if param is None:
            self.param = None
        else:
            self.param = np.asarray(param)
            if len(self.param) != n+1:
                raise TypeError("len(param) != n+1")
    #
    def init_param(self, random=True):
        if random:
            p0 = 2*np.random.random(self.n_param)-1
        else:
            p0 = np.zeros(self.n_param, 'd')
        
        if self.param is None:
            self.param = p0
        else:
            if len(p0) != len(self.param):
                raise TypeError("len(param) != len(p0)")
            self.param[:] = p0
    #
    def evaluate(self, Xk):
        return self.param[0] + Xk @ self.param[1:]
    #
    def gradient(self, Xk):
        G = np_empty(self.n_param, 'd')
        G[0] = 1.
        G[1:] = Xk
        return G
    #
    def gradient_x(self, Xk):
        return self.param[1:].copy()
    #
    def gradient_all(self, X):
        G = np_empty((len(X), self.n_param), 'd')
        G[:,0] = 1.0
        G[:,1:] = X
        return G


        